/*
*	Copyright(c) 2020 lutianming email：641471957@qq.com
*
*	Sheeps may be copied only under the terms of the GNU Affero General Public License v3.0
*/

#include "framework.h"
#include "ServerSocksProxy.h"
#include "ServerHook.h"
#include "ServerCase.h"
#include <algorithm>
#include <string>
#ifdef __WINDOWS__
#include "ws2tcpip.h"
#endif // __WINDOWS__

static int ProxyRemoteRequest(HSOCKET hsock, ServerProtocol* proto, const char* data, int len)
{
	h_SOCKS5_PROXY socks_info = proto->socks5_proxy;

	int hook_ret = 0;
	if (proxy_hook && server_hook) {
		hook_ret = server_hook->proxy_connect_send_hook("socks5", proto->initSock, socks_info->proxySock, TCP_PROTOCOL, data, len);
	}

	if (hook_ret == 0) HsocketSend(socks_info->proxySock, data, len);
	if (Proxy_record == true){
		insert_msg_to_record_list(socks_info->serverip, socks_info->serverport, socks_info->sessionid, 1, socks_info->ssl_conn ? SSL_PROTOCOL : TCP_PROTOCOL, data, len);
	}
	return len;
}

static int ProxyConnectRequest(HSOCKET hsock, ServerProtocol* proto, const char* data, int len)
{
	if (len < 10) return 0;
	//uint8_t ver = *data;
	uint8_t cmd = *(data + 1);
	uint8_t atype = *(data + 3);

	h_SOCKS5_PROXY socks_info = proto->socks5_proxy;

	if (cmd == 0x1) socks_info->proxytype = TCP_PROTOCOL;
	else if (cmd == 0x3) socks_info->proxytype = UDP_PROTOCOL;

	uint16_t port = ntohs(*((uint16_t*)(data + len - 2)));
	char ip[40] = { 0x0 };
	switch (atype){
	case 0x1:
		inet_ntop(AF_INET, (struct in_addr*)(data+4), ip, sizeof(ip));
		break;
	case 0x3:
		*(char*)(data + len - 2) = 0x0;
		GetHostByName(data + 5, ip, sizeof(ip));
		break;
	case 0x4:
		inet_ntop(AF_INET, (struct in6_addr*)(data + 4), ip, sizeof(ip));
		break;
	default:
		break;
	}
	socks_info->proxyStat = PROXY_CONNECTING;
	if (socks_info->proxytype == TCP_PROTOCOL)
	{
		char fromip[40] = { 0x0 };
		int fromport = 0;
		HsocketPeerAddr(hsock, fromip, sizeof(fromip), &fromport);
		LOG(slogid, LOG_FAULT, "%s:%d %p from[%s:%d] to[%s:%d]\n", __func__, __LINE__, proto, fromip, fromport, ip, port);

		char portstr[8] = { 0x0 };
		int portlen = snprintf(portstr, sizeof(portstr), "%d", port);
		const char* ports = config_get_string_value("server", "proxy_ssl_ports", NULL);
		if (ports) {
			const char* p_port = ports;
			const char* split = NULL;
			while (1) {
				split = strchr(p_port, ',');
				if (split) {
					if (split - p_port == portlen && strncmp(p_port, portstr, portlen) == 0) {
						socks_info->ssl_conn = true;
						break;
					}
					p_port = split + 1;
				}
				else {
					if (strncmp(p_port, portstr, portlen) == 0) socks_info->ssl_conn = true;
					break;
				}
			}
		}
		if (socks_info->ssl_conn) {
			socks_info->proxySock = HsocketConnect(proto, ip, port, SSL_PROTOCOL);
		}
		else {
			socks_info->proxySock = HsocketConnect(proto, ip, port, TCP_PROTOCOL);
		}
		if (socks_info->proxySock == NULL)
		{
			HsocketClose(proto->initSock);
		}
	}
	else if (socks_info->proxytype == UDP_PROTOCOL)
	{
		socks_info->udpClientSock = HsocketConnect(proto, "0.0.0.0", 0, UDP_PROTOCOL);
		ProxyConnectionMade(socks_info->udpClientSock, proto);
	}
	return len;
}

#define FD(x) x->fd

static int ProxyAuthRequest(HSOCKET hsock, ServerProtocol* proto, const char* data, int len)
{
	h_SOCKS5_PROXY socks_info = proto->socks5_proxy;
	int clen = *(data + 1);
	if (len < clen + 2)
		return 0;
	char buf[4] = { 0x0 };
	buf[0] = 0x5;
	buf[1] = 0x0;
	socks_info->proxyStat = PROXY_AUTHED;
	socks_info->sessionid = FD(hsock);
	HsocketSend(hsock, buf, 2);
	return clen + 2;
}

int CheckPoxyRequest(HSOCKET hsock, ServerProtocol* proto, const char* data, int len)
{
	h_SOCKS5_PROXY socks_info = proto->socks5_proxy;
	int clen = 0;
	if (hsock == proto->initSock)
	{
		switch (socks_info->proxyStat)
		{
		case PROXY_INIT:
			clen = ProxyAuthRequest(hsock, proto, data, len);
			break;
		case PROXY_AUTHED:
			clen = ProxyConnectRequest(hsock, proto, data, len);
			break;
		case PROXY_CONNECTED:
			clen = ProxyRemoteRequest(hsock, proto, data, len);
			break;
		default:
			break;
		}
	}
	else if (hsock == socks_info->proxySock && proto->initSock)
	{
		if (socks_info->proxytype == TCP_PROTOCOL)
		{
			int hook_ret = 0;
			if (proxy_hook && server_hook) {
				hook_ret = server_hook->proxy_connect_recv_hook("socks5", proto->initSock, socks_info->proxySock, TCP_PROTOCOL, data, len);
			}
			if (hook_ret == 0) HsocketSend(proto->initSock, data, len);
			if (Proxy_record)
				insert_msg_to_record_list(socks_info->serverip, socks_info->serverport, socks_info->sessionid, 3, socks_info->ssl_conn ? SSL_PROTOCOL : TCP_PROTOCOL, data, len);
		}
		else if (socks_info->proxytype == UDP_PROTOCOL)
		{
			char* buf = (char*)malloc((size_t)len + 10);
			if (buf)
			{
				memcpy(buf, socks_info->tempmsg, 10);
				memcpy(buf + 10, data, len);

				int hook_ret = 0;
				if (proxy_hook && server_hook) {
					hook_ret = server_hook->proxy_connect_recv_hook("socks5", socks_info->udpClientSock, socks_info->proxySock, TCP_PROTOCOL, data, len);
				}
				if (hook_ret == 0) HsocketSend(socks_info->udpClientSock, buf, len + 10);
				free(buf);
				if (Proxy_record)
					insert_msg_to_record_list(socks_info->serverip, socks_info->serverport, socks_info->sessionid, 3, UDP_PROTOCOL, data, len);
			}
			else
			{
				LOG(slogid, LOG_ERROR, "%s:%d malloc error\r\n", __func__, __LINE__);
				HsocketClose(proto->initSock);
			}
			clen = len;
		}
	}
	else if (hsock == socks_info->udpClientSock)
	{
		HsocketPeerAddr(hsock, socks_info->clientip, sizeof(socks_info->clientip), &socks_info->clientport);
		memcpy(socks_info->tempmsg, data, 10);

		inet_ntop(AF_INET, (data+4), socks_info->serverip, sizeof(socks_info->serverip));
		socks_info->serverport = ntohs(*(u_short*)(data + 8));
	
		if (socks_info->proxySock == NULL)
			socks_info->proxySock = HsocketConnect(proto, socks_info->serverip, socks_info->serverport, UDP_PROTOCOL);
		if (socks_info->proxySock == NULL) {
			printf("%s:%d create udp error\n", __func__, __LINE__);
		}

		int hook_ret = 0;
		if (proxy_hook && server_hook) {
			hook_ret = server_hook->proxy_connect_send_hook("socks5", socks_info->udpClientSock, socks_info->proxySock, TCP_PROTOCOL, data, len);
		}
		if (hook_ret == 0) {
			bool ret = HsocketSend(socks_info->proxySock, data + 10, len - 10);
			if (!ret) {
				printf("%s:%d send udp error\n", __func__, __LINE__);
			}
		}
		if (Proxy_record)
			insert_msg_to_record_list(socks_info->serverip, socks_info->serverport, socks_info->sessionid, 1, UDP_PROTOCOL, data + 10, len - 10);
		clen = len;
	}
	return clen;
}

void ProxyConnectionMade(HSOCKET hsock, ServerProtocol* proto)
{
	h_SOCKS5_PROXY socks_info = proto->socks5_proxy;
	if (socks_info->proxyStat == PROXY_CONNECTING && proto->initSock)
	{
		char buf[64] = { 0x0 };
		buf[0] = 0x5;
		buf[3] = 0x1;

		char local_ip[40] = { 0x0 };
		int local_port = 0;
		HsocketLocalAddr(proto->initSock, local_ip, sizeof(local_ip), &local_port);

		in_addr local_addr = { 0x0 };
		inet_pton(AF_INET, local_ip, &local_addr);
		memcpy(buf + 4, &local_addr, sizeof(local_addr));
			
		if (socks_info->proxytype == UDP_PROTOCOL){
			struct sockaddr_in6 addr_udp = {0x0};
			socklen_t nSize = sizeof(addr_udp);
			getsockname(socks_info->udpClientSock->fd, (struct sockaddr*)&addr_udp, &nSize);
			memcpy(buf + 4 + sizeof(in_addr), (char*)&addr_udp.sin6_port, 2);
		}
		else {
			unsigned short tcp_port = htons(local_port);
			memcpy(buf + 4 + sizeof(in_addr), (char*)&tcp_port, 2);
		}
		socks_info->proxyStat = PROXY_CONNECTED;
		HsocketSend(proto->initSock, buf, 4 + sizeof(in_addr) + 2);
		if (socks_info->ssl_conn)
			HsocketSSLCreate(proto->initSock, SSL_SERVER, 0, NULL, user_crt, pri_key);

		HsocketPeerAddr(hsock, socks_info->serverip, sizeof(socks_info->serverip), &socks_info->serverport);

		char proxy_addr[256] = { 0x0 };
		snprintf(proxy_addr, sizeof(proxy_addr), "从%s:%d  到%s:%d", socks_info->clientip, socks_info->clientport, socks_info->serverip, socks_info->serverport);
		ProxyAddr->insert(std::make_pair(proto, proxy_addr));

		if (proxy_hook && server_hook) {
			server_hook->proxy_connect_open_hook("socks5", "", proto->initSock, socks_info->proxySock, TCP_PROTOCOL);
		}

		if (Proxy_record == true && socks_info->proxytype == TCP_PROTOCOL){
			insert_msg_to_record_list(socks_info->serverip, socks_info->serverport, socks_info->sessionid, 0, socks_info->ssl_conn?SSL_PROTOCOL:TCP_PROTOCOL, NULL, 0);
		}
	}
	else if (hsock == proto->initSock) {
		LOG(slogid, LOG_FAULT, "%s:%d socks5 proxy upto ssl", __func__, __LINE__);
	}
	else{
		HsocketClose(hsock);
	}
}

void ProxyConnectionFailed(HSOCKET hsock, ServerProtocol* proto)
{
	LOG(slogid, LOG_DEBUG, "%s:%d\r\n", __func__, __LINE__);
	h_SOCKS5_PROXY socks_info = proto->socks5_proxy;
	if (hsock == socks_info->proxySock){
		if (socks_info->proxyStat == PROXY_CONNECTING && socks_info->retry > 0 && proto->initSock != NULL){
			socks_info->retry -= 1;
			char ip[40] = { 0x0 };
			int port = 0;
			HsocketPeerAddr(hsock, ip, sizeof(ip), &port);

			if (socks_info->ssl_conn) {
				socks_info->proxySock = HsocketConnect(proto, ip, port, SSL_PROTOCOL);
			}
			else {
				socks_info->proxySock = HsocketConnect(proto, ip, port, TCP_PROTOCOL);
			}
			if (socks_info->proxySock == NULL){
				HsocketClose(proto->initSock);
			}
		}
		else{
			if (proxy_hook && server_hook) {
				server_hook->proxy_connect_faile_hook("socks5", "", proto->initSock, socks_info->proxySock, TCP_PROTOCOL);
			}
			socks_info->proxySock = NULL;
			if (proto->initSock)HsocketClose(proto->initSock);
		}
	}
	else if(hsock == proto->initSock){
		proto->initSock = NULL;
		if (socks_info->proxySock) HsocketClose(socks_info->proxySock);
	}
}

void ProxyConnectionClosed(HSOCKET hsock, ServerProtocol* proto){
	h_SOCKS5_PROXY socks_info = proto->socks5_proxy;

	if (proxy_hook && server_hook && proto->initSock && socks_info->proxySock && socks_info->proxytype == TCP_PROTOCOL) {
		server_hook->proxy_connect_close_hook("socks5", proto->initSock, socks_info->proxySock, TCP_PROTOCOL);
	}

	if (hsock == socks_info->proxySock){
		socks_info->proxySock = NULL;
		HsocketClose(proto->initSock);
		HsocketClose(socks_info->udpClientSock);
	}
	else if (hsock == proto->initSock){
		//ProxyAddr->erase(proto);
		
		proto->initSock = NULL;
		if (socks_info->proxyStat >= PROXY_CONNECTING)
			HsocketClose(socks_info->proxySock);
		HsocketClose(socks_info->udpClientSock);
	}
	else if (hsock == socks_info->udpClientSock){
		socks_info->udpClientSock = NULL;
		HsocketClose(proto->initSock);
		HsocketClose(socks_info->proxySock);
	}
	if (socks_info->proxySock == NULL && proto->initSock == NULL && socks_info->udpClientSock == NULL){
		ProxyAddr->erase(proto);
		if (Proxy_record == true && socks_info->proxyStat >= PROXY_CONNECTED && socks_info->proxytype == TCP_PROTOCOL){
			insert_msg_to_record_list(socks_info->serverip, socks_info->serverport, socks_info->sessionid, 2, socks_info->ssl_conn ? SSL_PROTOCOL : TCP_PROTOCOL, NULL, 0);
		}
	}
}
